import argparse
from pathlib import Path
import shutil
import numpy as np
from transforms3d.quaternions import quat2mat
from tqdm import tqdm
from ..utils import to_pose, plan_move_qpos, execute_plan
from ..sim.wrapper_env import WrapperEnvConfig, WrapperEnv

generate_data_dir = Path("data")

num_data = 1000
start_index = 0


def parse_argument():
    parser = argparse.ArgumentParser(description="Launcher config - Physics")
    parser.add_argument("--robot", type=str, default='galbot')
    parser.add_argument("--obj", type=str, default='power_drill')
    parser.add_argument("--ctrl_dt", type=float, default=0.02)
    parser.add_argument("--headless", type=int, default=0)
    parser.add_argument("--reset_wait_steps", type=int, default=100)
    parser.add_argument("--num_data", type=int, default=1000)
    parser.add_argument("--start_index", type=int, default=0)

    args = parser.parse_args()
    global num_data, start_index
    num_data = args.num_data
    start_index = args.start_index
    return args

def construct_env(args:argparse.Namespace):
    env_config = WrapperEnvConfig(
        humanoid_robot=args.robot,
        obj_name=args.obj,
        headless=args.headless,
        ctrl_dt=args.ctrl_dt,
        reset_wait_steps=args.reset_wait_steps,
    )

    env = WrapperEnv(env_config)
    return env


def get_random_table_config():
    trans = np.array([0.55, 0.38, 0.68]) + np.random.rand(3) * np.array([0.05, 0.1, 0.06])
    size = np.array([0.68, 0.36, 0.02]) + np.random.rand(3) * np.array([0.04, 0.06, 0])
    return trans, size

def get_random_obj_config():
    obj_trans = np.array([0.5, 0.3, 0.82])
    obj_rot = np.random.randn(4)
    obj_rot /= np.linalg.norm(obj_rot)
    return obj_trans, obj_rot

def generate_data(env:WrapperEnv, save_path:str):

    table_trans, table_size = get_random_table_config()
    obj_trans, obj_rot = get_random_obj_config()
    env.set_table_obj_config(
        table_pose=to_pose(trans=table_trans),
        table_size=table_size,
        obj_pose=to_pose(trans=obj_trans, rot=quat2mat(obj_rot))
    )
    env.launch()
    env.reset(humanoid_qpos=np.array(
        [
            -1.32695362, -1.18605075, -0.0450398, -1.75000105, -1.34581286, -0.26273597, -1.10648651,
            0, 0, 0, 0, 0, 0
        ]
    ))
    for _ in range(100):
        env.step_env()



    obs_wrist = env.get_obs(camera_id=1)
    env.debug_save_obs(obs_wrist, save_path)

    driller_pose = env.get_driller_pose()
    if driller_pose[2, 3] < 0.5:
        env.close()
        return False

    with open(Path(save_path / "driller_pose.npy").as_posix(), "wb") as f:
        np.save(f, driller_pose)

    camera_intrinsic = env.sim.humanoid_robot_cfg.camera_cfg[1].intrinsics
    with open(Path(save_path / "camera_intrinsic.npy").as_posix(), "wb") as f:
        np.save(f, camera_intrinsic)

    env.close()
    return True


def main():
    args = parse_argument()
    env = construct_env(args)

    for i in tqdm(range(num_data)):
        save_path = generate_data_dir / f"{i + start_index:04d}"
        save_path.mkdir(parents=True, exist_ok=True)
        while generate_data(env, save_path) == False:
            pass


    env.close()

if __name__ == "__main__":
    main()